//! 32-bits Memory arena for types implementing `Copy`.
//! This Memory arena has been implemented to fit the use of tantivy's indexer
//! and has *twisted specifications*.
//!
//! - It works on stable rust.
//! - One can get an accurate figure of the memory usage of the arena.
//! - Allocation are very cheap.
//! - Allocation happening consecutively are very likely to have great locality.
//! - Addresses (`Addr`) are 32bits.
//! - Dropping the whole `MemoryArena` is cheap.
//!
//! # Limitations
//!
//! - Your object shall not implement `Drop`.
//! - `Addr` to the `Arena` are 32-bits. The maximum capacity of the arena is 4GB. *(Tantivy's
//!   indexer uses one arena per indexing thread.)*
//! - The arena only works for objects much smaller than  `1MB`. Allocating more than `1MB` at a
//!   time will result in a panic, and allocating a lot of large object (> 500KB) will result in a
//!   fragmentation.
//! - Your objects are store in an unaligned fashion. For this reason, the API does not let you
//!   access them as references.
//!
//! Instead, you store and access your data via `.write(...)` and `.read(...)`, which under the hood
//! stores your object using `ptr::write_unaligned` and `ptr::read_unaligned`.
use std::sync::Arc;
use std::{mem, ptr};

use once_cell::sync::Lazy;
use parking_lot::Mutex;

// Tanty's default is 20 bits, or 1MB.  Tantivy uses the memory arena during indexing and generally
// assumes long-running indexing threads.  We (pg_search) do indexing on the main thread and our
// indexing patterns are typically 1 document per segment.  :(  Using half tantivy's default memory
// saves quite a bit of indexing overhead
const NUM_BITS_PAGE_ADDR: usize = 19;
const PAGE_SIZE: usize = 1 << NUM_BITS_PAGE_ADDR; // pages are 512k large

// We use 32-bits addresses.
// - 19 bits for the in-page addressing
// - 13 bits for the page id.
// This limits us to 2^13 - 1=8191 for the page id.
const MAX_PAGES: usize = 1 << (32 - NUM_BITS_PAGE_ADDR);

/// Represents a pointer into the `MemoryArena`
/// .
/// Pointer are 32-bits and are split into
/// two parts.
///
/// The first 12 bits represent the id of a
/// page of memory.
///
/// The last 20 bits are an address within this page of memory.
#[derive(Copy, Clone, Debug)]
pub struct Addr(u32);

impl Addr {
    /// Creates a null pointer.
    #[inline]
    pub fn null_pointer() -> Addr {
        Addr(u32::MAX)
    }

    /// Returns the `Addr` object for `addr + offset`
    #[inline]
    pub fn offset(self, offset: u32) -> Addr {
        Addr(self.0.wrapping_add(offset))
    }

    #[inline]
    fn new(page_id: usize, local_addr: usize) -> Addr {
        Addr((page_id << NUM_BITS_PAGE_ADDR | local_addr) as u32)
    }

    #[inline]
    fn page_id(self) -> usize {
        (self.0 as usize) >> NUM_BITS_PAGE_ADDR
    }

    #[inline]
    fn page_local_addr(self) -> usize {
        (self.0 as usize) & (PAGE_SIZE - 1)
    }

    /// Returns true if and only if the `Addr` is null.
    #[inline]
    pub fn is_null(self) -> bool {
        self.0 == u32::MAX
    }
}

#[inline(always)]
pub fn store<Item: Copy + 'static>(dest: &mut [u8], val: Item) {
    debug_assert_eq!(dest.len(), std::mem::size_of::<Item>());
    unsafe {
        ptr::write_unaligned(dest.as_mut_ptr() as *mut Item, val);
    }
}

#[inline]
pub fn load<Item: Copy + 'static>(data: &[u8]) -> Item {
    debug_assert_eq!(data.len(), std::mem::size_of::<Item>());
    unsafe { ptr::read_unaligned(data.as_ptr() as *const Item) }
}

/// The `MemoryArena`
pub struct MemoryArena {
    pages: Vec<Page>,
    capacity: usize,
}

static ARENA_POOL: Lazy<Arc<Mutex<Vec<MemoryArena>>>> = Lazy::new(|| Default::default());

impl Default for MemoryArena {
    fn default() -> MemoryArena {
        ARENA_POOL.lock().pop().unwrap_or_else(|| {
            let first_page = Page::new(0);
            MemoryArena {
                pages: vec![first_page],
                capacity: 1,
            }
        })
    }
}

impl Drop for MemoryArena {
    fn drop(&mut self) {
        let my_pages = std::mem::replace(&mut self.pages, Vec::new());
        let my_capacity = self.capacity;
        let mut recycled = MemoryArena {
            pages: my_pages,
            capacity: my_capacity,
        };

        recycled.reset();
        ARENA_POOL.lock().push(recycled);
    }
}

impl MemoryArena {
    pub fn reset(&mut self) {
        unsafe {
            self.pages.set_len(1);
        }
        self.pages[0].len = 0;
    }

    /// Returns an estimate in number of bytes
    /// of resident memory consumed by the `MemoryArena`.
    ///
    /// Internally, it counts a number of `1MB` pages
    /// and therefore delivers an upperbound.
    pub fn mem_usage(&self) -> usize {
        self.pages.len() * PAGE_SIZE
    }

    /// Returns the number of bytes allocated in the arena.
    pub fn len(&self) -> usize {
        self.pages.len().saturating_sub(1) * PAGE_SIZE + self.pages.last().unwrap().len
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    #[inline]
    pub fn write_at<Item: Copy + 'static>(&mut self, addr: Addr, val: Item) {
        let dest = self.slice_mut(addr, std::mem::size_of::<Item>());
        store(dest, val);
    }

    /// Read an item in the memory arena at the given `address`.
    ///
    /// # Panics
    ///
    /// If the address is erroneous
    #[inline]
    pub fn read<Item: Copy + 'static>(&self, addr: Addr) -> Item {
        load(self.slice(addr, mem::size_of::<Item>()))
    }
    #[inline]
    fn get_page(&self, page_id: usize) -> &Page {
        unsafe { self.pages.get_unchecked(page_id) }
    }
    #[inline]
    fn get_page_mut(&mut self, page_id: usize) -> &mut Page {
        unsafe { self.pages.get_unchecked_mut(page_id) }
    }

    #[inline]
    pub fn slice(&self, addr: Addr, len: usize) -> &[u8] {
        self.get_page(addr.page_id())
            .slice(addr.page_local_addr(), len)
    }

    #[inline]
    pub fn slice_from(&self, addr: Addr) -> &[u8] {
        self.get_page(addr.page_id())
            .slice_from(addr.page_local_addr())
    }
    #[inline]
    pub fn slice_from_mut(&mut self, addr: Addr) -> &mut [u8] {
        self.get_page_mut(addr.page_id())
            .slice_from_mut(addr.page_local_addr())
    }

    #[inline]
    pub fn slice_mut(&mut self, addr: Addr, len: usize) -> &mut [u8] {
        self.get_page_mut(addr.page_id())
            .slice_mut(addr.page_local_addr(), len)
    }

    /// Add a page and allocate len on it.
    /// Return the address
    fn add_page(&mut self, len: usize) -> Addr {
        let npages = self.pages.len();

        if npages + 1 < self.capacity {
            // we have a hidden, pre-allocated page

            unsafe {
                // make it live
                self.pages.set_len(npages + 1);
                let page = self.pages.get_unchecked_mut(npages);
                page.len = len;
                Addr::new(page.page_id, 0)
            }
        } else {
            // must allocate a new page and add it to the arena
            let new_page_id = self.pages.len();
            let mut page = Page::new(new_page_id);
            page.len = len;
            self.pages.push(page);
            self.capacity += 1;
            Addr::new(new_page_id, 0)
        }
    }

    /// Allocates `len` bytes and returns the allocated address.
    #[inline]
    pub fn allocate_space(&mut self, len: usize) -> Addr {
        let page_id = self.pages.len() - 1;
        if let Some(addr) = self.get_page_mut(page_id).allocate_space(len) {
            return addr;
        }
        self.add_page(len)
    }
}

struct Page {
    page_id: usize,
    len: usize,
    data: Box<[u8; PAGE_SIZE]>,
}

impl Page {
    fn new(page_id: usize) -> Page {
        assert!(page_id < MAX_PAGES);
        Page {
            page_id,
            len: 0,
            data: vec![0u8; PAGE_SIZE].into_boxed_slice().try_into().unwrap(),
        }
    }

    #[inline]
    fn is_available(&self, len: usize) -> bool {
        len + self.len <= PAGE_SIZE
    }

    #[inline]
    fn slice(&self, local_addr: usize, len: usize) -> &[u8] {
        let data = &self.slice_from(local_addr);
        unsafe { data.get_unchecked(..len) }
    }

    #[inline]
    fn slice_from(&self, local_addr: usize) -> &[u8] {
        &self.data[local_addr..]
    }
    #[inline]
    fn slice_from_mut(&mut self, local_addr: usize) -> &mut [u8] {
        &mut self.data[local_addr..]
    }

    #[inline]
    fn slice_mut(&mut self, local_addr: usize, len: usize) -> &mut [u8] {
        let data = &mut self.data[local_addr..];
        unsafe { data.get_unchecked_mut(..len) }
    }

    #[inline]
    fn allocate_space(&mut self, len: usize) -> Option<Addr> {
        if self.is_available(len) {
            let addr = Addr::new(self.page_id, self.len);
            self.len += len;
            Some(addr)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {

    use super::{MemoryArena, MAX_PAGES};
    use crate::memory_arena::PAGE_SIZE;

    #[test]
    fn test_arena_allocate_slice() {
        let mut arena = MemoryArena::default();
        let a = b"hello";
        let b = b"happy tax payer";

        let addr_a = arena.allocate_space(a.len());
        arena.slice_mut(addr_a, a.len()).copy_from_slice(a);

        let addr_b = arena.allocate_space(b.len());
        arena.slice_mut(addr_b, b.len()).copy_from_slice(b);

        assert_eq!(arena.slice(addr_a, a.len()), a);
        assert_eq!(arena.slice(addr_b, b.len()), b);
    }

    #[test]
    fn test_arena_allocate_end_of_page() {
        let mut arena = MemoryArena::default();

        // A big block
        let len_a = PAGE_SIZE - 2;
        let addr_a = arena.allocate_space(len_a);
        *arena.slice_mut(addr_a, len_a).last_mut().unwrap() = 1;

        // Single bytes
        let addr_b = arena.allocate_space(1);
        arena.slice_mut(addr_b, 1)[0] = 2;

        let addr_c = arena.allocate_space(1);
        arena.slice_mut(addr_c, 1)[0] = 3;

        let addr_d = arena.allocate_space(1);
        arena.slice_mut(addr_d, 1)[0] = 4;

        assert_eq!(arena.slice(addr_a, len_a)[len_a - 1], 1);
        assert_eq!(arena.slice(addr_b, 1)[0], 2);
        assert_eq!(arena.slice(addr_c, 1)[0], 3);
        assert_eq!(arena.slice(addr_d, 1)[0], 4);
    }

    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
    struct MyTest {
        pub a: usize,
        pub b: u8,
        pub c: u32,
    }

    #[test]
    fn test_store_object() {
        let mut arena = MemoryArena::default();
        let a = MyTest {
            a: 143,
            b: 21,
            c: 32,
        };
        let b = MyTest {
            a: 113,
            b: 221,
            c: 12,
        };

        let num_bytes = std::mem::size_of::<MyTest>();
        let addr_a = arena.allocate_space(num_bytes);
        arena.write_at(addr_a, a);

        let addr_b = arena.allocate_space(num_bytes);
        arena.write_at(addr_b, b);

        assert_eq!(arena.read::<MyTest>(addr_a), a);
        assert_eq!(arena.read::<MyTest>(addr_b), b);
    }

    #[test]
    fn test_can_allocate_4GB() {
        use super::NUM_BITS_PAGE_ADDR;

        let mut arena = MemoryArena::default();
        for i in 0..MAX_PAGES {
            let addr = arena.allocate_space(PAGE_SIZE - 1); // -1 to ensure we don't cross page boundary
            assert_eq!(addr.page_id(), i);
        }
        assert_eq!(arena.pages.len(), 1 << (32 - NUM_BITS_PAGE_ADDR));
        assert_eq!(
            arena.mem_usage(),
            PAGE_SIZE * (1 << (32 - NUM_BITS_PAGE_ADDR))
        );
    }
    #[test]
    #[should_panic]
    fn test_cannot_allocate_more_than_4GB() {
        use super::NUM_BITS_PAGE_ADDR;

        let mut arena = MemoryArena::default();
        for i in 0..MAX_PAGES + 1 {
            let addr = arena.allocate_space(PAGE_SIZE - 1); // -1 to ensure we don't cross page boundary
            assert_eq!(addr.page_id(), i);
        }
        assert_eq!(arena.pages.len(), 1 << (32 - NUM_BITS_PAGE_ADDR));
        assert_eq!(
            arena.mem_usage(),
            PAGE_SIZE * (1 << (32 - NUM_BITS_PAGE_ADDR))
        );
    }
}
